import os
import sys
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import nibabel as nib

matplotlib.use("Agg")  # 无 GUI 环境下使用


def save_slices_as_pngs(nii_path, output_dir, num_slices=121):
    # 提取文件名（不带扩展名）
    base_name = os.path.splitext(os.path.basename(nii_path))[0]

    # 加载 .nii 文件
    img = nib.load(nii_path)
    data = img.get_fdata()

    # 创建横截面、冠状面和矢状面的子目录
    axial_dir = os.path.join(output_dir, "axial")
    coronal_dir = os.path.join(output_dir, "coronal")
    sagittal_dir = os.path.join(output_dir, "sagittal")
    os.makedirs(axial_dir, exist_ok=True)
    os.makedirs(coronal_dir, exist_ok=True)
    os.makedirs(sagittal_dir, exist_ok=True)

    # 保存横截面切片
    max_axial_slices = min(data.shape[2], num_slices)
    for i in range(max_axial_slices):
        slice_data = np.rot90(data[:, :, i], k=2)  # 旋转 180 度
        plt.imshow(slice_data.T, cmap='gray')
        plt.axis('off')
        output_path = os.path.join(axial_dir, f"{i + 1:03d}.png")
        print("保存路径:", output_path)
        plt.savefig(output_path, bbox_inches='tight', pad_inches=0)
        plt.close()

    # 保存冠状面切片
    max_coronal_slices = min(data.shape[1], num_slices)
    for i in range(max_coronal_slices):
        slice_data = np.rot90(data[:, i, :], k=2)  # 旋转 180 度
        plt.imshow(slice_data.T, cmap='gray')
        plt.axis('off')
        output_path = os.path.join(coronal_dir, f"{i + 1:03d}.png")
        print("保存路径:", output_path)
        plt.savefig(output_path, bbox_inches='tight', pad_inches=0)
        plt.close()

    # 保存矢状面切片
    max_sagittal_slices = min(data.shape[0], num_slices)
    for i in range(max_sagittal_slices):
        slice_data = np.rot90(data[i, :, :], k=2)  # 旋转 180 度
        plt.imshow(slice_data.T, cmap='gray')
        plt.axis('off')
        output_path = os.path.join(sagittal_dir, f"{i + 1:03d}.png")
        print("保存路径:", output_path)
        plt.savefig(output_path, bbox_inches='tight', pad_inches=0)
        plt.close()

    print(f"已成功保存横截面、冠状面和矢状面的 PNG 切片至目录: {output_dir}")


if __name__ == "__main__":
    nii_file = sys.argv[1]
    output_dir = sys.argv[2]
    # 调用转换函数
    save_slices_as_pngs(nii_file, output_dir)
    print(f"切片已保存到目录: {output_dir}")